package com.fw.leetcode.tree;

import com.fw.leetcode.LeetCode;

import java.util.ArrayDeque;

/**
 * 230. Kth Smallest Element in a BST
 *
 * Given the root of a binary search tree, and an integer k, return the kth smallest value (1-indexed)
 * of all the values of the nodes in the tree.
 *
 * Example 1:
 *  Input: root = [3,1,4,null,2], k = 1
 *  Output: 1
 *
 * Example 2:
 *  Input: root = [5,3,6,2,4,null,null,1], k = 3
 *  Output: 3
 *
 * Constraints:
 *  The number of nodes in the tree is n.
 *  1 <= k <= n <= 10^4
 *  0 <= Node.val <= 10^4
 *
 * Follow up: If the BST is modified often (i.e., we can do insert and delete operations) and
 * you need to find the kth smallest frequently, how would you optimize?
 */
public class Num_0230 implements LeetCode {
    private interface Solution extends Tree {
        int kthSmallest(TreeNode root, int k);

        default void assertExpected(Integer[] root, int k, int expected) {
            TreeNode rooNode = buildBinaryTree(root);
            if (kthSmallest(rooNode, k) != expected) {
                kthSmallest(rooNode, k);
            }
        }
    }

    private static class MySolution implements Solution {

        @Override
        public int kthSmallest(TreeNode root, int k) { // 最优：中序遍历(递归，主动停止)：时O(log n) 空O(log n)
            /*
             * 二叉搜索树（排序树）中第k小的元素，计数从1开始
             * 采用 [中序遍历] (递归)
             * 已优化点：找到第 k 个后能主动停止，不需要遍历整个树
             *
             *             5
             *        3         6
             *     2    4
             *  1
             */
            return kthSmallest(root, new TreeNode(k)); // 建立辅助节点传递 k 值递减信息
        }

        private int kthSmallest(TreeNode root, TreeNode k) {
            if (root == null) return -1;
            int left = kthSmallest(root.left, k); // 左
            if (left > -1) {
                return left; // 找到元素提前停止
            }
            if (k.val > 0 && --k.val <= 0) { // 中，倒计数减为0时，即是所求
                return root.val; // 找到第 k 个元素后，不再往后遍历，递归返回
            }
            return kthSmallest(root.right, k); // 右
        }
    }

    private static class MyStackSolution implements Solution {

        @Override
        public int kthSmallest(TreeNode root, int k) { // 最优：中序遍历(栈，主动停止)：时O(<n) 空O(<n)
            /*
             * 二叉搜索树（排序树）中第k小的元素，计数从1开始
             * 采用 [中序遍历] (栈)
             * 优点：找到第 k 个后可停止，不需要遍历整个树
             *
             *             5
             *        3         6
             *     2    4
             *  1
             */
            ArrayDeque<TreeNode> stack = new ArrayDeque<>();
            TreeNode cur = root;

            while (cur != null || !stack.isEmpty()) {
                while (cur != null) { // 左，直接遍历到最左子树叶节点
                    stack.push(cur);
                    cur = cur.left;
                }
                cur = stack.pop(); // 中，出栈访问根，若计数达到 k 则返回
                if (--k <= 0) {
                    return cur.val;
                }
                cur = cur.right; // 右
            }
            return -1;
        }
    }

    public static void main(String[] args) {
        Solution solution = new MyStackSolution();
        solution.assertExpected(new Integer[]{3,1,4,null,2}, 1, 1);
        solution.assertExpected(new Integer[]{5,3,6,2,4,null,null,1}, 3, 3);
    }
}
